import os
import base64
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding as asym_padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from OpenSSL import crypto

# ----------------- 新增CA证书加载 -----------------
def load_ca_cert(ca_path="ca.crt"):
    """加载CA根证书"""
    print_header("0. 加载CA根证书")
    try:
        with open(ca_path, "rb") as f:
            ca_cert = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
        
        # 打印CA证书信息
        print("[CA证书信息]")
        print(f"  主题: {ca_cert.get_subject().CN}")
        print(f"  颁发者: {ca_cert.get_issuer().CN}")
        print(f"  有效期: {ca_cert.get_notBefore().decode()} 至 {ca_cert.get_notAfter().decode()}")
        return ca_cert
        
    except FileNotFoundError:
        print(f"[错误] 找不到CA证书文件: {ca_path}")
        exit(1)
    except Exception as e:
        print(f"[错误] CA证书加载失败: {str(e)}")
        exit(1)

# ----------------- 修改后的验证流程 -----------------
# 加载虚拟机证书
def load_vm_cert(vm_path="vm.crt"):
    """加载虚拟机证书"""
    print_header("1. 加载虚拟机证书")
    try:
        with open(vm_path, "rb") as f:
            vm_cert = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
        
        print("[虚拟机证书信息]")
        print(f"  主题: {vm_cert.get_subject().CN}")
        print(f"  颁发者: {vm_cert.get_issuer().CN}")
        return vm_cert
        
    except FileNotFoundError:
        print(f"[错误] 找不到虚拟机证书文件: {vm_path}")
        exit(1)


# ----------------- 3. 数据验证和解密 -----------------
def verify_and_decrypt(session_key,vm_cert):
    """验证签名并解密数据"""
    print_header("3. 加载加密数据文件")
    
    # 加载加密数据文件
    try:
        with open("iv.bin", "rb") as f:
            iv = f.read()
        with open("tag.bin", "rb") as f:
            tag = f.read()
        with open("encrypted_data.bin", "rb") as f:
            encrypted_data = f.read()
        with open("signature.bin", "rb") as f:
            signature = f.read()
            
        print("[文件加载成功]")
        print(f"  IV长度: {len(iv)} bytes")
        print(f"  标签长度: {len(tag)} bytes")
        print(f"  加密数据长度: {len(encrypted_data)} bytes")
        print(f"  签名长度: {len(signature)} bytes")
        
    except FileNotFoundError as e:
        print(f"[错误] 缺少必要文件: {e.filename}")
        exit(1)

    # 组合待验证数据
    data_to_sign = iv + tag + encrypted_data
    print_hex(data_to_sign, "待验证数据摘要: ", length=20)

    # 加载虚拟机公钥
    vm_pubkey = vm_cert.get_pubkey().to_cryptography_key()
    print("\n[使用的公钥信息]")
    print(f"  密钥类型: {type(vm_pubkey).__name__}")
    print(f"  密钥长度: {vm_pubkey.key_size} bits")

    # 验证签名
    print_header("4. 验证数字签名")
    try:
        vm_pubkey.verify(
            signature,
            data_to_sign,
            asym_padding.PSS(
                mgf=asym_padding.MGF1(hashes.SHA256()),
                salt_length=asym_padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        print("[签名验证] 状态: 有效 (数据完整且来源可信)")
    except Exception as e:
        print(f"[签名验证] 状态: 无效! 原因: {type(e).__name__} - {str(e)}")
        exit(1)

    # 解密数据
    print_header("5. 解密数据")
    try:
        cipher = Cipher(
            algorithms.AES(session_key),
            modes.GCM(iv, tag)
        )
        decryptor = cipher.decryptor()
        decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()
        
        print("[解密成功]")
        print(f"  解密数据长度: {len(decrypted_data)} bytes")
        print_hex(decrypted_data, "原始数据: ", length=32)
        print(f"  UTF-8可读内容: {decrypted_data[:64].decode('utf-8', 'replace')}...")
        
        return decrypted_data
    except Exception as e:
        print(f"[解密失败] 错误类型: {type(e).__name__}, 详情: {str(e)}")
        exit(1)


# ----------------- 主验证流程 -----------------
def main_verify():
    # 加载证书
    ca_cert = load_ca_cert()  # 新增的关键加载步骤
    vm_cert = load_vm_cert()
    # 验证证书链
    print_header("2. 验证证书有效性")
    store = crypto.X509Store()
    store.add_cert(ca_cert)
    store_ctx = crypto.X509StoreContext(store, vm_cert)
    
    try:
        store_ctx.verify_certificate()
        print("[证书验证] 状态: 有效 (由CA签发)")
    except crypto.X509StoreContextError as e:
        print(f"[证书验证] 状态: 无效! 原因: {str(e)}")
        exit(1)
    
    # ----------------- 新增RSA解密流程 -----------------
    print_header("2.5 解密会话密钥")
    try:
        # 加载虚拟机私钥
        with open("vm.key", "rb") as f:
            vm_private_key = serialization.load_pem_private_key(
                f.read(),
                password=None
            )
        
        # 读取加密的会话密钥
        with open("encrypted_session_key.bin", "rb") as f:
            encrypted_session_key = f.read()
        
        # RSA解密得到原始AES密钥
        session_key = vm_private_key.decrypt(
            encrypted_session_key,
            asym_padding.OAEP(
                mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        
        print_hex(session_key, "解密后的AES密钥: ", length=8)
        
        # 验证密钥长度
        if len(session_key) not in [16, 24, 32]:
            print(f"[错误] 无效的AES密钥长度: {len(session_key)} bytes")
            exit(1)
            
    except Exception as e:
        print(f"[解密失败] 错误类型: {type(e).__name__}, 详情: {str(e)}")
        exit(1)
    
    
    # 执行验证和解密
    decrypted_data = verify_and_decrypt(session_key,vm_cert)
    
    print("\n" + "="*40)
    print("=== 完整验证流程完成 ===")
    print("="*40)
    return decrypted_data

# --------------- 辅助函数 ---------------
def print_header(title):
    print(f"\n{'='*40}")
    print(f"=== {title}")
    print(f"{'='*40}\n")

def print_hex(data, prefix="", length=16):
    """格式化打印二进制数据的十六进制值（截断长数据）"""
    hex_str = data.hex()
    truncated = hex_str[:length*2]
    if len(hex_str) > length*2:
        truncated += "..."
    print(f"{prefix}{truncated} [Total {len(data)} bytes]")



if __name__ == "__main__":
    decrypted = main_verify()
